#!/bin/bash

# 快速JupyterHub用户管理脚本
# 直接使用JupyterHub API添加用户，无需重启或重载配置

set -e

# 配置
JUPYTER_CONFIG_FILE="jupyterhub_config.py"
JUPYTER_CONFIG_BACKUP="jupyterhub_config.py.backup"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 显示帮助信息
show_help() {
    echo "快速JupyterHub用户管理脚本"
    echo ""
    echo "用法: $0 <命令> [参数]"
    echo ""
    echo "命令:"
    echo "  add-users-fast <container_name> <user_list_file> <api_token>  - 快速批量添加用户（推荐）"
    echo "  add-user-fast <container_name> <username> <api_token>         - 快速添加单个用户（推荐）"
    echo "  add-users-config <container_name> <user_list_file>            - 通过配置重载添加用户（传统方式）"
    echo "  health-check <container_name>                                  - 检查服务健康状态"
    echo "  get-api-token <container_name>                                 - 获取API Token"
    echo "  list-users <container_name>                                    - 列出所有用户"
    echo "  remove-user <container_name> <username>                        - 删除用户"
    echo ""
    echo "示例:"
    echo "  $0 add-users-fast jhub-test-group-001 users.txt abc123"
    echo "  $0 add-user-fast jhub-test-group-001 student001 abc123"
    echo "  $0 health-check jhub-test-group-001"
    echo ""
    echo "用户列表文件格式 (每行一个用户名):"
    echo "  student001"
    echo "  student002"
    echo "  student003"
}

# 检查容器是否存在
check_container() {
    local container_name=$1
    
    if ! docker ps | grep -q "$container_name"; then
        log_error "容器 $container_name 不存在或未运行"
        return 1
    fi
    
    log_info "容器 $container_name 状态正常"
    return 0
}

# 获取容器端口
get_container_port() {
    local container_name=$1
    local port=$(docker port "$container_name" 8000/tcp | cut -d: -f2)
    
    if [ -z "$port" ]; then
        log_error "无法获取容器端口映射"
        return 1
    fi
    
    echo "$port"
}

# 检查JupyterHub服务健康状态
health_check() {
    local container_name=$1
    local port=$(get_container_port "$container_name")
    
    local health_url="http://127.0.0.1:$port/hub/health"
    log_info "检查健康状态: $health_url"
    
    if curl -s -f "$health_url" > /dev/null; then
        log_success "JupyterHub服务健康"
        return 0
    else
        log_error "JupyterHub服务不健康"
        return 1
    fi
}

# 获取API Token
get_api_token() {
    local container_name=$1
    local group_id=${container_name#jhub-}
    local config_dir="/tmp/practice-platform/shared/${group_id}-jhub"
    local config_file="$config_dir/$JUPYTER_CONFIG_FILE"
    
    if [ ! -f "$config_file" ]; then
        log_error "配置文件不存在: $config_file"
        return 1
    fi
    
    # 从配置文件中提取API Token
    local api_token=$(grep "api_token" "$config_file" | sed "s/.*'api_token': '\([^']*\)'.*/\1/")
    
    if [ -n "$api_token" ]; then
        echo "$api_token"
        return 0
    else
        log_error "配置文件中未找到API Token"
        return 1
    fi
}

# 快速添加用户（推荐方式）
add_user_fast() {
    local container_name=$1
    local username=$2
    local api_token=$3
    
    log_info "快速添加用户: $username 到容器: $container_name"
    
    # 检查容器状态
    if ! check_container "$container_name"; then
        return 1
    fi
    
    # 检查服务健康状态
    if ! health_check "$container_name"; then
        log_error "JupyterHub服务不健康，无法添加用户"
        return 1
    fi
    
    # 获取端口
    local port=$(get_container_port "$container_name")
    
    # 通过API添加用户
    local api_url="http://127.0.0.1:$port/jhub/hub/api/users/$username"
    log_info "调用API添加用户: $api_url"
    
    local response=$(curl -s -w "%{http_code}" -X POST \
        -H "Authorization: token $api_token" \
        -H "Content-Type: application/json" \
        "$api_url")
    
    local http_code="${response: -3}"
    local response_body="${response%???}"
    
    if [ "$http_code" = "201" ] || [ "$http_code" = "200" ]; then
        log_success "用户 $username 添加成功 (HTTP $http_code)"
        return 0
    else
        log_error "用户 $username 添加失败 (HTTP $http_code): $response_body"
        return 1
    fi
}

# 快速批量添加用户（推荐方式）
add_users_fast() {
    local container_name=$1
    local user_list_file=$2
    local api_token=$3
    
    log_info "快速批量添加用户到容器: $container_name"
    
    # 检查容器状态
    if ! check_container "$container_name"; then
        return 1
    fi
    
    # 检查服务健康状态
    if ! health_check "$container_name"; then
        log_error "JupyterHub服务不健康，无法添加用户"
        return 1
    fi
    
    if [ ! -f "$user_list_file" ]; then
        log_error "用户列表文件不存在: $user_list_file"
        return 1
    fi
    
    # 读取用户列表并添加
    local success_count=0
    local total_count=0
    
    while IFS= read -r username; do
        username=$(echo "$username" | tr -d '[:space:]')
        if [ -n "$username" ]; then
            total_count=$((total_count + 1))
            log_info "添加用户 $total_count: $username"
            
            if add_user_fast "$container_name" "$username" "$api_token"; then
                success_count=$((success_count + 1))
            fi
            
            # 短暂延迟，避免API调用过于频繁
            sleep 0.5
        fi
    done < "$user_list_file"
    
    log_success "批量添加用户完成: 成功 $success_count/$total_count"
    return $([ $success_count -eq $total_count ] && echo 0 || echo 1)
}

# 通过配置重载添加用户（传统方式）
add_users_config() {
    local container_name=$1
    local user_list_file=$2
    
    log_info "通过配置重载添加用户到容器: $container_name"
    
    # 检查容器状态
    if ! check_container "$container_name"; then
        return 1
    fi
    
    if [ ! -f "$user_list_file" ]; then
        log_error "用户列表文件不存在: $user_list_file"
        return 1
    fi
    
    local group_id=${container_name#jhub-}
    local config_dir="/tmp/practice-platform/shared/${group_id}-jhub"
    local config_file="$config_dir/$JUPYTER_CONFIG_FILE"
    
    if [ ! -f "$config_file" ]; then
        log_error "JupyterHub配置文件不存在: $config_file"
        return 1
    fi
    
    # 备份配置文件
    cp "$config_file" "$config_dir/$JUPYTER_CONFIG_BACKUP"
    log_info "配置文件已备份"
    
    # 读取现有用户列表
    local existing_users=""
    if grep -q "c.Authenticator.allowed_users" "$config_file"; then
        existing_users=$(grep "c.Authenticator.allowed_users" "$config_file" | sed "s/.*allowed_users = {\([^}]*\)}.*/\1/")
    fi
    
    # 读取新用户列表
    local new_users=""
    while IFS= read -r username; do
        username=$(echo "$username" | tr -d '[:space:]')
        if [ -n "$username" ]; then
            if [ -n "$new_users" ]; then
                new_users="$new_users, '$username'"
            else
                new_users="'$username'"
            fi
        fi
    done < "$user_list_file"
    
    # 合并用户列表
    local all_users=""
    if [ -n "$existing_users" ] && [ -n "$new_users" ]; then
        all_users="$existing_users, $new_users"
    elif [ -n "$existing_users" ]; then
        all_users="$existing_users"
    elif [ -n "$new_users" ]; then
        all_users="$new_users"
    fi
    
    # 更新配置文件
    if [ -n "$all_users" ]; then
        sed -i "s/c\.Authenticator\.allowed_users = {[^}]*}/c.Authenticator.allowed_users = {$all_users}/" "$config_file"
        log_info "配置文件已更新"
        
        # 尝试重载配置
        log_info "尝试重载配置..."
        if docker exec "$container_name" pkill -HUP -f jupyterhub; then
            log_info "SIGHUP信号已发送"
            sleep 5
            
            if health_check "$container_name"; then
                log_success "配置重载成功"
                return 0
            else
                log_warning "配置重载可能失败，但用户已添加到配置文件"
                return 1
            fi
        else
            log_warning "发送SIGHUP信号失败，但用户已添加到配置文件"
            return 1
        fi
    else
        log_warning "没有用户需要添加"
        return 0
    fi
}

# 列出所有用户
list_users() {
    local container_name=$1
    local group_id=${container_name#jhub-}
    local config_dir="/tmp/practice-platform/shared/${group_id}-jhub"
    local config_file="$config_dir/$JUPYTER_CONFIG_FILE"
    
    log_info "列出容器 $container_name 的所有用户"
    
    if [ ! -f "$config_file" ]; then
        log_error "配置文件不存在: $config_file"
        return 1
    fi
    
    if grep -q "c.Authenticator.allowed_users" "$config_file"; then
        local users=$(grep "c.Authenticator.allowed_users" "$config_file" | sed "s/.*allowed_users = {\([^}]*\)}.*/\1/")
        if [ -n "$users" ]; then
            log_info "当前用户列表:"
            echo "$users" | tr ',' '\n' | sed 's/^[[:space:]]*//; s/[[:space:]]*$//' | grep -v '^$'
        else
            log_info "当前没有配置用户"
        fi
    else
        log_info "配置文件中没有找到用户列表"
    fi
}

# 删除用户
remove_user() {
    local container_name=$1
    local username=$2
    local group_id=${container_name#jhub-}
    local config_dir="/tmp/practice-platform/shared/${group_id}-jhub"
    local config_file="$config_dir/$JUPYTER_CONFIG_FILE"
    
    log_info "删除用户: $username 从容器: $container_name"
    
    # 检查容器状态
    if ! check_container "$container_name"; then
        return 1
    fi
    
    # 备份配置文件
    cp "$config_file" "$config_dir/$JUPYTER_CONFIG_BACKUP"
    log_info "配置文件已备份"
    
    # 从配置文件中删除用户
    if sed -i "s/'$username',\?//g; s/,\s*}/}/g; s/{,\s*}/{}/g" "$config_file"; then
        log_success "用户 $username 已从配置文件删除"
        
        # 尝试重载配置
        log_info "尝试重载配置..."
        if docker exec "$container_name" pkill -HUP -f jupyterhub; then
            log_info "SIGHUP信号已发送"
            sleep 5
            
            if health_check "$container_name"; then
                log_success "配置重载成功"
                return 0
            else
                log_warning "配置重载可能失败，但用户已从配置文件删除"
                return 1
            fi
        else
            log_warning "发送SIGHUP信号失败，但用户已从配置文件删除"
            return 1
        fi
    else
        log_error "删除用户失败"
        return 1
    fi
}

# 主函数
main() {
    if [ $# -eq 0 ]; then
        show_help
        exit 1
    fi
    
    local command=$1
    shift
    
    case "$command" in
        "add-users-fast")
            if [ $# -ne 3 ]; then
                log_error "用法: $0 add-users-fast <container_name> <user_list_file> <api_token>"
                exit 1
            fi
            add_users_fast "$1" "$2" "$3"
            ;;
        "add-user-fast")
            if [ $# -ne 3 ]; then
                log_error "用法: $0 add-user-fast <container_name> <username> <api_token>"
                exit 1
            fi
            add_user_fast "$1" "$2" "$3"
            ;;
        "add-users-config")
            if [ $# -ne 2 ]; then
                log_error "用法: $0 add-users-config <container_name> <user_list_file>"
                exit 1
            fi
            add_users_config "$1" "$2"
            ;;
        "health-check")
            if [ $# -ne 1 ]; then
                log_error "用法: $0 health-check <container_name>"
                exit 1
            fi
            health_check "$1"
            ;;
        "get-api-token")
            if [ $# -ne 1 ]; then
                log_error "用法: $0 get-api-token <container_name>"
                exit 1
            fi
            get_api_token "$1"
            ;;
        "list-users")
            if [ $# -ne 1 ]; then
                log_error "用法: $0 list-users <container_name>"
                exit 1
            fi
            list_users "$1"
            ;;
        "remove-user")
            if [ $# -ne 2 ]; then
                log_error "用法: $0 remove-user <container_name> <username>"
                exit 1
            fi
            remove_user "$1" "$2"
            ;;
        "help"|"-h"|"--help")
            show_help
            ;;
        *)
            log_error "未知命令: $command"
            show_help
            exit 1
            ;;
    esac
}

# 执行主函数
main "$@" 